package main

import (
    "encoding/base64"
    "fmt"
    "github.com/gliderlabs/ssh"
    gossh "golang.org/x/crypto/ssh"
    "io"
    "os"
    "syscall"
    "unsafe"
    "os/exec"
    "github.com/kr/pty"
    "strings"
)

type SshServer struct {
    Port    string
    PemPath string
    auth    Auth
    srv     *ssh.Server
}

func NewSshServer(port, pem, secrets, keys string) (srv *SshServer, err error) {
    if _, err = os.Stat(pem); err != nil {
        return
    }
    if _, err = os.Stat(secrets); err != nil {
        return
    }
    if _, err = os.Stat(keys); err != nil {
        return
    }
    auth, err := NewLocalAuth(secrets, keys)
    if err != nil {
        return
    }
    return &SshServer{
        Port:    port,
        PemPath: pem,
        auth:    auth,
    }, nil
}

func setWinsize(f *os.File, w, h int) { //调整term窗口大小使用 不需要动
    syscall.Syscall(syscall.SYS_IOCTL, f.Fd(), uintptr(syscall.TIOCSWINSZ),
        uintptr(unsafe.Pointer(&struct{ h, w, x, y uint16 }{uint16(h), uint16(w), 0, 0})))
}

func chop(addr string) string {
    if strings.Contains(addr, ":") {
        s := strings.Split(addr, ":")
        return s[0]
    }
    return addr
}

func (ss *SshServer) publicKeyAuth(ctx ssh.Context, key ssh.PublicKey) bool { //使用公钥登陆的验证
    addr := chop(ctx.RemoteAddr().String())
    user := ctx.User()
    k := base64.StdEncoding.EncodeToString(gossh.MarshalAuthorizedKey(key))
    ctx.SetValue("x-public-key", k)
    ok := ss.auth.AuthKey(user, addr, k)
    log.Infof("%s try to login with key, %v", user, ok)
    return ok
}

func (ss *SshServer) passwdAuth(ctx ssh.Context, password string) bool { //使用密码登陆的验证
    user := ctx.User()
    ok := ss.auth.AuthUserPass(user, password, false)
    log.Infof("%s try to login with pass, %v", user, ok)
    return ok
}

func (ss *SshServer) serve() error {
    ssh.Handle(ss.handle)
    ss.srv = &ssh.Server{Addr: ss.Port, Handler: nil}
    ss.srv.SetOption(ssh.HostKeyFile(ss.PemPath))
    ss.srv.SetOption(ssh.PublicKeyAuth(ss.publicKeyAuth)) //使用密码登陆的设置
    ss.srv.SetOption(ssh.PasswordAuth(ss.passwdAuth))     //使用密码登陆的设置
    log.Infof("start ssh server")
    return ss.srv.ListenAndServe()
}

func (ss *SshServer) handle(session ssh.Session) {
    input := session.Command()
    addr := chop(session.RemoteAddr().String())
    user := session.User()
    log.Infof("handle request %s from %s", user, addr)
    switch {
    case len(input) == 1 && input[0] == "add-key":
        log.Infof("handle %s request %s from %s", "add key", user, addr)
        k := session.Context().Value("x-public-key")
        if k == nil {
            io.WriteString(session, "key is null\n")
            log.Infof("request err is: %s, %s from %s", "key is null", user, addr)
            session.Exit(2)
        }
        key := k.(string)
        if err := ss.auth.AddKey(user, addr, key); err != nil {
            io.WriteString(session, "\n"+err.Error()+"\n")
            log.Infof("request err is: %v, %s from %s", err, user, addr)
            session.Exit(2)
        }
        if err := ss.auth.Unload(); err != nil {
            io.WriteString(session, "\n"+err.Error()+"\n")
            log.Infof("request err is: %v, %s from %s", err, user, addr)
            session.Exit(2)
        }
        io.WriteString(session, "added\n")
        log.Infof("request result is %s, %s from %s", "added", user, addr)
        session.Exit(0)
    case len(input) == 2 && input[0] == "change-pass-to":
        log.Infof("handle %s request %s from %s", "change pass", user, addr)
        if err := ss.auth.ChangePassTo(user, input[1]); err != nil {
            io.WriteString(session, fmt.Sprintf("change pass error: %v\n", err))
            log.Infof("request err is: %v, %s from %s", err, user, addr)
            session.Exit(2)
        } else {
            io.WriteString(session, "changed\n")
            log.Infof("request result is %s, %s from %s", "changed", user, addr)
            session.Exit(0)
        }
    default:
        log.Info("handle %s request %s from %s", "bash", user, addr)
        ptyReq, winCh, isPty := session.Pty()
        if !isPty {
            io.WriteString(session, "No PTY requested\n try to use:\n ssh -tt xx.xx.xx.xx command")
            session.Exit(2)
        }
        cmd := exec.Command("bash")
        cmd.Env = append(cmd.Env, os.Environ()...)
        cmd.Env = append(cmd.Env, fmt.Sprintf("TERM=%s", ptyReq.Term))
        f, err := pty.Start(cmd)
        if err != nil {
            io.WriteString(session, fmt.Sprintf("command error: %v", err))
            log.Infof("request err is: %v, %s from %s", err, user, addr)
            session.Exit(2)
        }
        go func() { //调整窗口大小
            for win := range winCh {
                setWinsize(f, win.Width, win.Height)
            }
        }()
        go func() { //读输入
            io.Copy(f, session) // stdin
        }()
        io.Copy(session, f)
        return
    }
}
